#include"head.h"
#include"common.h"
#include"conn.h"
static struct IPRule *ipRuleHead = NULL;
static DEFINE_RWLOCK(ipRuleLock);
void int_to_ip(unsigned int ip_int, char *ip_str) {
    // 使用位运算提取每个字节
    uint8_t byte1 = (ip_int >> 24) & 0xFF;
    uint8_t byte2 = (ip_int >> 16) & 0xFF;
    uint8_t byte3 = (ip_int >> 8) & 0xFF;
    uint8_t byte4 = ip_int & 0xFF;

    // 生成字符串格式的 IP 地址
    sprintf(ip_str, "%u.%u.%u.%u", byte1, byte2, byte3, byte4);
}

// 在名称为after的规则后新增一条规则，after为空时则在首部新增一条规则
struct IPRule * addIPRuleToChain(struct IPRule rule) {
    struct IPRule *newRule,*now;
    newRule = (struct IPRule *) kzalloc(sizeof(struct IPRule), GFP_KERNEL);
    if(newRule == NULL) {
        printk(KERN_WARNING "[fw iprule.c] kzalloc fail.\n");
        return NULL;
    }
    memcpy(newRule, &rule, sizeof(struct IPRule));
    // 新增规则至规则链表
    write_lock(&ipRuleLock);
    if(rule.action != NF_ACCEPT) 
         eraseConnRelated(rule); // 消除新增规则的影响

    if(ipRuleHead == NULL) {
        ipRuleHead = newRule;
        ipRuleHead->nx = NULL;
        ipRuleHead->pre = NULL;
        write_unlock(&ipRuleLock);
        return newRule;
    }
    // if(strlen(after)==0) {
    //     newRule->nx = ipRuleHead;
    //     ipRuleHead = newRule;
    //     write_unlock(&ipRuleLock);
    //     return newRule;
    // }
  
    for(now=ipRuleHead;now->nx!=NULL;now=now->nx) ;

       
            newRule->nx = NULL;
            newRule->pre=now;
            now->nx = newRule;
            write_unlock(&ipRuleLock);
            return newRule;
        
    
    // 添加失败
    write_unlock(&ipRuleLock);
    kfree(newRule);
    return NULL;
}

// 删除所有名称为name的规则
int delIPRuleFromChain(char name[]) {
    struct IPRule *now,*tmp;
  int count=0;
  
    write_lock(&ipRuleLock);
    if(ipRuleHead!=NULL && strcmp(ipRuleHead->name,name)==0) 
    {
        tmp = ipRuleHead;
        ipRuleHead = ipRuleHead->nx;
        if(ipRuleHead!=NULL) ipRuleHead->pre=NULL;
        eraseConnRelated(*tmp); // 消除删除规则的影响
        kfree(tmp);
      count=1;
    }
   else{
     for(now=ipRuleHead;now!=NULL && now->nx!=NULL;) {
        if(strcmp(now->nx->name,name)==0)
         { 
            tmp = now->nx;
            
            now->nx->nx->pre=now;

               now->nx = now->nx->nx; // 更新当前节点的下一个指针
                if (now->nx != NULL) {
                    now->nx->pre = now; // 如果下一个节点不为空，则更新它的前驱指针
                }
             eraseConnRelated(*tmp); // 消除删除规则的影响
            kfree(tmp);
           count++;
           break;
        } 
        else 
        {
            now = now->nx;
        }
    }

   }
    write_unlock(&ipRuleLock);
   
    return count;
}

// 将所有规则形成Netlink回包
void* formAllIPRules(unsigned int *len) {
    struct KernelResponse *head;
    struct IPRule *now;
    void *mem,*p;
    unsigned int count;
    read_lock(&ipRuleLock);
    for(now=ipRuleHead,count=0;now!=NULL;now=now->nx,count++);
    *len = sizeof(struct KernelResponse) + sizeof(struct IPRule)*count;

    mem = kzalloc(*len, GFP_ATOMIC);
    if(mem == NULL) {
        printk(KERN_WARNING "[fw rules] kzalloc fail.\n");
        read_unlock(&ipRuleLock);
        return NULL;
    }

    head = (struct KernelResponse *)mem;
  
    // head->rspLen = count;
    head->rsplen =sizeof(struct IPRule)*count;
    head->arrlen=count;
    for(now=ipRuleHead,p=(mem + sizeof(struct KernelResponse));now!=NULL;now=now->nx,p=p+sizeof(struct IPRule))
        
        memcpy(p, now, sizeof(struct IPRule));
    
    
    read_unlock(&ipRuleLock);

    return mem;
}


bool isIPMatch(unsigned int ipl, unsigned int ipr, unsigned int mask) {
	return (ipl & mask) == (ipr & mask);
}
void getPort(struct sk_buff *skb, struct iphdr *hdr, unsigned short *src_port, unsigned short *dst_port){
	struct tcphdr *tcpHeader;
	struct udphdr *udpHeader;
	switch(hdr->protocol){
		case IPPROTO_TCP:
			//printk("TCP protocol\n");
			tcpHeader = (struct tcphdr *)(skb->data + (hdr->ihl * 4));
			*src_port = ntohs(tcpHeader->source);
			*dst_port = ntohs(tcpHeader->dest);
			break;
		case IPPROTO_UDP:
			//printk("UDP protocol\n");
			udpHeader = (struct udphdr *)(skb->data + (hdr->ihl * 4));
			*src_port = ntohs(udpHeader->source);
			*dst_port = ntohs(udpHeader->dest);
			break;
		case IPPROTO_ICMP:
		default:
			//printk("other protocol\n");
			*src_port = 0;
			*dst_port = 0;
			break;
	}
}
bool matchOneRule(struct IPRule *rule,
 unsigned int sip, unsigned int dip, unsigned short sport, unsigned int dport, u_int8_t proto) {
    return (isIPMatch(sip,rule->saddr,rule->smask) &&
			isIPMatch(dip,rule->daddr,rule->dmask) &&
			(sport >= ((unsigned short)(rule->sport >> 16)) && sport <= ((unsigned short)(rule->sport & 0xFFFFu))) &&
			(dport >= ((unsigned short)(rule->dport >> 16)) && dport <= ((unsigned short)(rule->dport & 0xFFFFu))) &&
			(rule->protocol == IPPROTO_IP || rule->protocol == proto));
}
struct IPRule matchIPRules(struct sk_buff *skb, int *isMatch) {
    struct IPRule *now,ret;
	unsigned short sport,dport;
	struct iphdr *header = ip_hdr(skb);
	*isMatch = 0;
	getPort(skb,header,&sport,&dport);
	read_lock(&ipRuleLock);
	for(now=ipRuleHead;now!=NULL;now=now->nx) {
		if(matchOneRule(now,ntohl(header->saddr),ntohl(header->daddr),sport,dport,header->protocol)) {
				ret = *now;
				*isMatch = 1;
				break;
		}
	}
	read_unlock(&ipRuleLock);
	return ret;
}


// // 进行过滤规则匹配，isMatch存储是否匹配到规则
// struct IPRule matchIPRules(struct sk_buff *skb, int *isMatch) {
//     struct IPRule *now,ret;
// 	unsigned short sport,dport;
// 	struct iphdr *header = ip_hdr(skb);
// 	*isMatch = 0;
// 	getPort(skb,header,&sport,&dport);
// 	read_lock(&ipRuleLock);
// 	for(now=ipRuleHead;now!=NULL;now=now->nx) {
// 		if(matchOneRule(now,ntohl(header->saddr),ntohl(header->daddr),sport,dport,header->protocol)) {
// 				ret = *now;
// 				*isMatch = 1;
// 				break;
// 		}
// 	}
// 	read_unlock(&ipRuleLock);
// 	return ret;
// }
